package main

import (
	"flag"
	"fmt"
	"http"
	"log"
	"net"
	"os"
	"strings"
	"strconv"
	"websocket"
	"encoding/base64"
)

var origin = flag.String("origin",
	"",
	"Specify trusted origin, for example, http://ssh.example.org. Use 'any' to allow any origin.")

var targetPort = flag.Int("target_port",
	-1,
	"Specify target tcp port allowed to connect to. Use 0 to allow all ports.")

var defaultTargetPort = flag.Int("default_target_port",
	22,
	"Default tcp port if it's not specified")

var serverAddr string
var maxSize = 10 * 1024 * 1024

func Reverse(c net.Conn, ws *websocket.Conn) {
	buf := make([]byte, maxSize)
	enc_buf := make([]byte, base64.StdEncoding.EncodedLen(maxSize))
	for {
		// log.Print("Entering c.Read...")
		n, err := c.Read(buf)
		if err != nil {
			log.Print("c.Read failed: ", err)
			err = ws.Close()
			if err != nil {
				log.Print("ws.Close failed: ", err)
			}
			return
		}
		// log.Print("Bytes received from real socket:", n)
		base64.StdEncoding.Encode(enc_buf, buf[0:n])
		enc_len := base64.StdEncoding.EncodedLen(n)
		ws.Write(enc_buf[0:enc_len])
		// log.Print("Bytes sent to websocket:", enc_len)
	}
}

func echoServer(ws *websocket.Conn) {
	if *origin != "any" && ws.Origin != *origin {
		log.Printf("Incoming request came from %s. Blocked. Trusted origin is %s", ws.Origin, *origin)
		return
	}
	log.Printf("Incoming request received from %s", ws.Origin)
	buf := make([]byte, 1000)
	n, err := ws.Read(buf)
	if err != nil {
		log.Print("Error: ", err)
		return
	}
	hostname := string(buf[0:n])
	port := *defaultTargetPort
	switch strings.Count(hostname, ":") {
	case 0:
	case 1:
		a := strings.Split(hostname, ":", 2)
		var err os.Error
		port, err = strconv.Atoi(a[1])
		if err == nil {
			hostname = a[0]
			break
		}
		fallthrough
	default:
		log.Printf("Can't parse hostname:port from %s", hostname)
		return
	}
	if *targetPort != 0 && *targetPort != port {
		log.Printf("Port %d is not allowed. Port %d is allowed", port, *targetPort)
		return
	}
	str := fmt.Sprintf("%s:%d", hostname, port)
	log.Print("Incoming msg: ", str)
	c, err := net.Dial("tcp", "", str)
	if err != nil {
		log.Print("net.Dial failed: ", err)
		return
	}
	log.Print("Connected to " + str)
	go Reverse(c, ws)

	base64_reader := base64.NewDecoder(base64.StdEncoding, ws)
	for {
		// log.Print("Entering websocket.Read...")
		n, err := base64_reader.Read(buf)
		if err != nil {
			log.Print("ws.Read failed: ", err)
			err = c.Close()
			if err != nil {
				log.Print("c.Close failed", err)
			}
			return
		}
		// log.Print("Data received from websocket: ", n)
		_, err = c.Write(buf[0:n])
		if err != nil {
			log.Print("c.Write failed: ", err)
		}
	}
}


func main() {
	flag.Parse()
	if *origin == "" {
		flag.PrintDefaults()
		log.Fatal("--origin must be specified")
	}
	if *targetPort == -1 {
		log.Fatal("--target_port must be specified")
	}
	l, e := net.Listen("tcp", ":10101") // any available address
	if e != nil {
		log.Fatalf("net.Listen tcp :0 %v", e)
	}
	serverAddr = l.Addr().String()
	log.Print("Test WebSocket server listening on ", serverAddr)
	http.Handle("/echo", websocket.Handler(echoServer))
	http.Handle("/echoDraft75", websocket.Draft75Handler(echoServer))
	http.Serve(l, nil)
}
